#!/usr/bin/env python
# util.py

"""
Authors: Sandra Smit (sandra.smit@colorado.edu) and 
    Greg Caporaso (caporaso@colorado.edu)

Status: Development

Revision History:
6/1-7/13/04: Inital development and testing
4/8/05 Sandra Smit: changed parameter setting (SuppressStderr etc) in the
__init__ of CommandLineApplication, so that passing in 'False' actually works 
7/28/05 Rob Knight: added _input_as_multiline_string so we can efficiently
dump in large strings without additional processing.

12/14/05: Jeremy Widmann: mkstemp() creates a tuple: (file_descriptor,path).
    The file descriptor is an open file, while the path is a unique path to be
    used for a temporary file.  The file descriptor remains open during the
    python session if not handled.  As a result calling mkstemp() many times and
    not closing the file descriptor will cause you to run out of memory.
    Changed all places where mkstemp is called to get the file descriptor out
    of the tuple and close it. 
3/13/06 Micah Hamady: replaced mkstemp with self.getTmpFilename, fixed bug with
    _input_filename. Also now cleans up temp input_filename.

Conventions for docstrings for app controllers:
1. Name of the application, and a brief high-level overview of what that
   application does. For example, for BLAST we would say that it takes a
   set of sequences and the name of a database, and returns information
   about the matching sequences in the database.
1. URL where the application can be downloaded, or statement that it is not
   available.
2. URL for the documentation, if available, or statement that it is not
   available.
3. Version(s) of program that the app controller works for.


"""

import commands
from sys import platform
from os import remove,system,mkdir,getcwd,close
from old_cogent.app.parameters import Parameter, FlagParameter, ValuedParameter,\
    MixedParameter,Parameters, _find_synonym, is_not_None
from old_cogent.util.misc import if_
from random import choice
 
#the following are used to create temp file names       
_chars = "abcdefghigklmnopqrstuvwxyz"
_all_chars = _chars + _chars.upper() + "0123456790"


class ApplicationError(OSError):
    pass
    
class ResultPath(object):
    """ Hold a file path a boolean value specifying whether file was written
    """
    def __init__(self,Path,IsWritten=True):
        """ Initialize the ResultPath object

            Path: a string representing the absolute or relative path where
                the file can be found
            IsWritten: a boolean specifying whether the file has been written,
                default = True
        """
        self.Path = Path
        self.IsWritten = IsWritten

class CommandLineAppResult(dict):
    """ Class for holding the result of a CommandLineApplication run """

    def __init__(self,out,err,exit_status,result_paths):
        """Initialization of CommandLineAppResult

        out: a file handler to the file containing the stdout
        err: a file handler to the file containing the stderr
        exit_status: the exit status of the program, 0 if run ok, 1 else.
        result_paths: dictionary containing ResultPath objects for each 
            output file that could be written
        input_handler_temp_files: set of filenames automatically generated by 
            one of the input handlers
        """
        
        self['StdOut'] = out
        self['StdErr'] = err
        self['ExitStatus'] = exit_status
       
        self.file_keys = result_paths.keys()
        for key,value in result_paths.items():
            if value.IsWritten:
                try:
                    self[key] = open(value.Path)
                except IOError:
                    raise ApplicationError, 'Could not open %s' %value.Path
            else:
                self[key] = None

    def cleanUp(self):
        """ Delete files that are written by CommandLineApplication from disk
            
            WARNING: after cleanUp() you may still have access to part of 
                your result data, but you should be aware that if the file
                size exceeds the size of the buffer you will only have part 
                of the file. To be safe, you should not use cleanUp() until 
                you are done with the file or have copied it to a different 
                location.
        """
        file_keys = self.file_keys
        for item in file_keys:
            if self[item] is not None:
                self[item].close()
                remove(self[item].name)

        # remove input handler temp files
        if hasattr(self, "_input_filename"): 
            remove(self._input_filename)

    def __del__(self):
        """ Delete temporary files created by the CommandLineApplication 
        """
        if self['StdOut'] is not None:
            remove(self['StdOut'].name)
        if self['StdErr'] is not None:
            remove(self['StdErr'].name)

class Application(object):
    """ Generic Class for controlling an application """

    _command = None
    _command_delimiter = ' '
    _parameters = {}
    _synonyms = {}
    
    def __init__(self,params=None):
        """
            params: a dict of parameters which should be turned on where the 
                key is either the parameter id or a synonym for the parameter
                and the value is either the value for the parameter or None
        """
        self.Parameters = Parameters(self._parameters, self._synonyms)
        if params:
            for key,v in params.items():
                try:
                    self.Parameters[key].on(v)
                except TypeError:
                    self.Parameters[key].on()

class CommandLineApplication(Application):
    """ Generic class for controlling command line applications 
    """

    _input_handler = '_input_as_string'
    _suppress_stderr = False
    _suppress_stdout = False
    _working_dir = None
    TmpPrefix = 'tmp'
    TmpSuffix = '.txt'

    def __init__(self,params=None,InputHandler=None,SuppressStderr=None,\
        SuppressStdout=None,WorkingDir=None,TmpDir='/tmp', \
        TmpNameLen=20, HALT_EXEC=False):
        """ Initialize the CommandLineApplication object
        
            params: a dictionary mapping the Parameter id or synonym to its
                value (or None for FlagParameters or MixedParameters in flag
                mode) for Parameters that should be turned on
            InputHandler: this is the method to be run on data when it is
                passed into call. This should be a string containing the
                method name. The default is _input_as_string which casts data
                to a string before appending it to the command line argument
            SuppressStderr: if set to True, will route standard error to
                /dev/null, False by default
            SuppressStdout: if set to True, will route standard out to
                /dev/null, False by default
            WorkingDir: the directory where you want the application to run,
                default is the current working directory, but is useful to 
                change in cases where the program being run creates output
                to its current working directory and you either don't want
                it to end up where you are running the program, or the user 
                running the script doesn't have write access to the current 
                working directory
                WARNING: WorkingDir MUST be an absolute path!
            TmpDir: the directory where temp files will be created, /tmp
                by default 
            TmpNameLen: the length of the temp file name
            HALT_EXEC: if True, raises exception w/ command output just
            before execution, doesn't clean up temp files. Default False.
        """
        # set attributes to parameter that was passed in or class default
        if InputHandler is not None:
            self.InputHandler = InputHandler
        else:
            self.InputHandler = self._input_handler
        if SuppressStderr is not None:
            self.SuppressStderr = SuppressStderr
        else:
            self.SuppressStderr = self._suppress_stderr
        if SuppressStdout is not None:
            self.SuppressStdout = SuppressStdout
        else:
            self.SuppressStdout = self._suppress_stdout
        if WorkingDir is not None:
            self.WorkingDir = WorkingDir
        else:
            self.WorkingDir = self._working_dir or getcwd()
        self.TmpDir = TmpDir
        self.TmpNameLen = TmpNameLen
        self.HaltExec = HALT_EXEC
        #===========================
        #try: 
        #    mkdir(self.WorkingDir)
        #except OSError:
            # Directory already exists
        #    pass 
        #===========================
        # create a variable to hold the name of the file being used as
        # input to the application. this is important especially when 
        # you are using an input handler which creates a temporary file
        # and the output filenames are based on the input filenames
        self._input_filename = None
        
        super(CommandLineApplication,self).__init__(params=params)
    
    def __call__(self,data=None):
        """Run the application with the specified kwargs on data
        
            data: anything that can be cast into a string or written out to
                a file. Usually either a list of things or a single string or 
                number. input_handler will be called on this data before it 
                is passed as part of the command-line argument, so by creating
                your own input handlers you can customize what kind of data
                you want you application to accept
        """
        input_handler = self.InputHandler
        suppress_stdout = self.SuppressStdout
        suppress_stderr = self.SuppressStderr
        if suppress_stdout:
            outfile = '/dev/null'
        else:
            outfile = self.getTmpFilename(self.WorkingDir)
        if suppress_stderr:
            errfile = '/dev/null'
        else:
            errfile = self.getTmpFilename(self.WorkingDir)
        if data is None:
            input_arg = ''
        else:
            input_arg = getattr(self,input_handler)(data)

        # Build up the command, consisting of a BaseCommand followed by
        # input and output (file) specifications
        command = self._command_delimiter.join(filter(None,\
            [self.BaseCommand,input_arg,'>',outfile,'2>',errfile]))
   
        # DEBUG
        #if 'clust' in command:
        #print command
        #    raise ValueError, command
        
        if self.HaltExec: 
            raise AssertionError, "Halted exec with command:\n" + command
        # The return value of system is a 16-bit number containing the signal 
        # number that killed the process, and then the exit status. 
        # We only want to keep the exit status so do a right bitwise shift to 
        # get rid of the signal number byte
        exit_status = system(command) >> 8
      
        # Determine if error should be raised due to exit status of 
        # appliciation
        if not self._accept_exit_status(exit_status):
            raise ApplicationError, \
             'Unacceptable application exit status: %s, command: %s'\
                % (str(exit_status),command)
        
        # open the stdout and stderr if not being suppressed
        out = None
        if not suppress_stdout:
            out = open(outfile,"r")
        err = None        
        if not suppress_stderr:
            err = open(errfile,"r")
       
        result =  CommandLineAppResult(out,err,exit_status,\
            result_paths=self._get_result_paths(data)) 

        # Clean up the input file if one was created
        if self._input_filename:
            remove(self._input_filename)
            self._input_filename = None

        return result
   
    def _input_as_string(self,data):
        """ Return data as a string """
        return str(data)

    def _input_as_multiline_string(self, data):
        """Write a multiline string to a temp file and return the filename.

        data: a multiline string to be written to a file.
        """
        filename = self._input_filename = self.getTmpFilename(self.WorkingDir)
        data_file = open(filename,'w')
        data_file.write(data)
        data_file.close()
        return filename
   
    def _input_as_lines(self,data):
        """ Write a seq of lines to a temp file and return the filename string
        
            data: a sequence to be written to a file, each element of the 
                sequence will compose a line in the file

            Note: '\n' will be stripped off the end of each sequence element
                before writing to a file in order to avoid multiple new lines
                accidentally be written to a file
        """
        filename = self._input_filename = self.getTmpFilename(self.WorkingDir)
        data_file = open(filename,'w')
        data_to_file = '\n'.join([str(d).strip('\n') for d in data])
        data_file.write(data_to_file)
        data_file.close()
        return filename
   
    def _get_base_command(self):
        """ Returns the full command string 

            input_arg: the argument to the command which represents the input 
                to the program, this will be a string, either 
                representing input or a filename to get input from
        """
        command_parts = []
        # Append a change directory to the beginning of the command to change 
        # to self.WorkingDir before running the command
        cd_command = ''.join(['cd ',self.WorkingDir,';'])
        if self._command is None:
            raise ApplicationError, '_command has not been set.'
        command = self._command
        parameters = self.Parameters
        synonyms = self._synonyms
        
        command_parts.append(cd_command)
        command_parts.append(command)
        command_parts.append(self._command_delimiter.join(filter(\
            None,(map(str,parameters.values())))))
      
        return self._command_delimiter.join(command_parts).strip()
    
    BaseCommand = property(_get_base_command)
    
    def _get_WorkingDir(self):
        """Gets the working directory"""
        return self._curr_working_dir

    def _set_WorkingDir(self,path):
        """Sets the working directory
        
        Appends a slash to the end of path
        The reasoning behind this is that the user may or may not pass
        in a path with a '/' at the end. Since having multiple
        '/' at the end doesn't hurt anything, it's convienient to 
        be able to rely on it, and not have to check for it
        """
        self._curr_working_dir = path + '/'
        try:
            mkdir(self.WorkingDir)
        except OSError:
            # Directory already exists
            pass 

    WorkingDir = property(_get_WorkingDir,_set_WorkingDir)

    
    def _accept_exit_status(self,exit_status):
        """ Return False to raise an error due to exit_status of applciation
        
            This method should be overwritten if you'd like to raise an error
            based on certain exit statuses of the application that was run. The
            default is that no value of exit_status will raise an error.
        """
        return True

    def _get_result_paths(self,data):
        """ Return a dict of ResultPath objects representing all possible output
            
            This method should be overwritten if the application creates output
            other than stdout and stderr.  This dictionary will have keys based
            on the name that you'd like to access the file by in the 
            CommandLineAppResult object that will be created, and the values
            which are ResultPath objects. For an example of how this should be
            written see the rnaview or vienna_package classes.
            WARNING: be sure that the path that you give a file is accurate
                from any directory where the program could be running. For that
                reason, absolute paths are very good. Relative paths can also be
                used as long as you are careful. For cases where the application
                leaves files in the current working directory, you should append
                self.WorkingDir to the beginning of the file name.
                It would be a very bad idea to just use a file name as the path,
                in some cases that you might not be testing for.
        """
        return {}

    def getTmpFilename(self, tmp_dir="/tmp"):
        # temp hack - change this to lookup and generate file in class

        # check not none
        if not tmp_dir:
            tmp_dir = self.TmpDir
        # if not current directory, append "/" if not already on path
        elif not tmp_dir.endswith("/"):
            tmp_dir += "/"

        return ''.join([tmp_dir, self.TmpPrefix,
            ''.join([choice(_all_chars) for i in range(self.TmpNameLen)]),
            self.TmpSuffix])

def get_tmp_filename(tmp_dir="/tmp", prefix="tmp"):
    """
    Generate temp filename
    """
    # check not none
    if not tmp_dir:
        tmp_dir = ""
    # if not current directory, append "/" if not already on path
    elif not tmp_dir.endswith("/"):
        tmp_dir += "/"

    chars = "abcdefghigklmnopqrstuvwxyz"
    picks = chars + chars.upper() + "0123456790"
    return tmp_dir + prefix + "%s.txt" % ''.join([choice(picks) for i in range(20)]) 


